package a4ny;

import java.util.ArrayList;
import java.util.Set;
import java.util.TreeMap;
import java.lang.Math;

public class CorpusProcesser {
	TreeMap<String, TreeMap<String, Integer>> wordCounts;
	TreeMap<String, Integer> posCounts;
	TreeMap<Bigram, Integer> bigramCounts;
	
	TreeMap<String, TreeMap<String, Double>> wordProb;
	TreeMap<String, Double> posProb;
	TreeMap<Bigram, Double> bigramProb;
	
	public CorpusProcesser(){
		wordCounts = new TreeMap<String, TreeMap<String, Integer>>();
		posCounts = new TreeMap<String, Integer>();
		bigramCounts = new TreeMap<Bigram, Integer>();
		
		wordProb = new TreeMap<String,TreeMap<String, Double> >();
		posProb = new TreeMap<String, Double>();
		bigramProb = new TreeMap<Bigram, Double>();
	}
	
	public void process(ArrayList<String> words, ArrayList<String> pos){
		TreeMap<String, Integer> tagCount;
		for(int i = 0; i < words.size(); i++){
			// Extract distinct words and frequencies
			String thisWord = words.get(i);
			String thisPos = pos.get(i);
			if(wordCounts.containsKey(thisWord)){
				tagCount = wordCounts.get(thisWord);
				int count1 = tagCount.containsKey(thisPos) ? tagCount.get(thisPos) : 0;
				tagCount.put(thisPos, count1 + 1);
				wordCounts.put(thisWord, tagCount);
			}else{
				tagCount = new TreeMap<String, Integer>();
				tagCount.put(thisPos, 1);
				wordCounts.put(thisWord, tagCount);
			}

			// Extract distinct pos and frequencies
			int count2 = posCounts.containsKey(thisPos) ? posCounts.get(thisPos) : 0;
			posCounts.put(thisPos, count2 + 1);
			posCounts.put("<s>", 1);
			
			// Extract distinct pos bigrams and frequencies
			Bigram b;
			if(i != 0){
				b = new Bigram(thisPos, pos.get(i-1));
			}else{
				b = new Bigram(thisPos, "<s>");
			}
			int count3 = bigramCounts.containsKey(b) ? bigramCounts.get(b) : 0;
			bigramCounts.put(b, count3 + 1);
		}
		
		// Word probabilities
		TreeMap<String, Integer> countPerPos;
		TreeMap<String, Double> tagProb;
		Set<String> tags;
		Set<String> wordsUnique = wordCounts.keySet();
		for(String w : wordsUnique){
			countPerPos = wordCounts.get(w);
			tags = countPerPos.keySet();
			tagProb = new TreeMap<String, Double>();
			for(String t : tags){
				tagProb.put(t, (double)countPerPos.get(t)/(double)posCounts.get(t));
			}
			wordProb.put(w, tagProb);
		}
		
		// Pos probabilities
		Set<String> posUnique = posCounts.keySet();
		int totalNbrPos = 0;
		for(String p : posUnique){
			totalNbrPos += posCounts.get(p);
		}
		for(String p : posUnique){
			double d = (double)posCounts.get(p)/(double)totalNbrPos;
			posProb.put(p, d);
		}

		
		// Bigram probabilities
		Set<Bigram> bigramsUnique = bigramCounts.keySet();
		double totalPosProb = 0.0;
		for(Bigram b : bigramsUnique){
			double d = (double)bigramCounts.get(b)/(double)posCounts.get(b.condition);
			bigramProb.put(b, d);
			totalPosProb += d;
		}
		System.out.println(totalPosProb);
	}
	
	public TreeMap<String, TreeMap<String, Integer> > getWordCounts(){
		return wordCounts;
	}
	
	public TreeMap<String, Integer> getPosCounts(){
		return posCounts;
	}
	
	public TreeMap<Bigram, Integer> getBigramCounts(){
		return bigramCounts;
	}
	
	public TreeMap<String, TreeMap<String, Double> > getWordProb(){
		return wordProb;
	}
	
	public TreeMap<String, Double> getPosProb(){
		return posProb;
	}
	
	public TreeMap<Bigram, Double> getBigramProb(){
		return bigramProb;
	}
}
